package com.spiro.test.mr;

import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.FieldVisitor;
import org.objectweb.asm.MethodVisitor;

import static org.objectweb.asm.Opcodes.*;

/**
 * @Author: Shaoping Huang
 * @Description:
 * @Date: 1/3/2018
 */
public class JobAdapter extends ClassVisitor {

    public static final String JOB_FIELD_NAME = "currentJob";
    public static final String JOB_FIELD_DESC = "Lorg/apache/hadoop/mapreduce/Job;";

    private String owner;
    private boolean isInterface;

    public JobAdapter(ClassVisitor cv) {
        super(ASM5, cv);
    }

    @Override
    public void visit(int version, int access, String name, String signature,
                      String superName, String[] interfaces) {
        super.visit(version, access, name, signature, superName, interfaces);
        owner = name;
        isInterface = (access & ACC_INTERFACE) == 1;
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);

        if (isInterface || mv == null) {
            return mv;
        }

        if (access == 0
                && name.equals("<init>")
                && desc.equals("(Lorg/apache/hadoop/mapred/JobConf;)V")) {
            return new JobMethodAdapter(mv);
        }

        return mv;
    }

    @Override
    public void visitEnd() {
        if (!isInterface) {
            FieldVisitor fv = super.visitField(
                    (ACC_PUBLIC | ACC_STATIC),
                    JOB_FIELD_NAME,
                    JOB_FIELD_DESC,
                    null, null);
            if (fv != null) {
                fv.visitEnd();
            }
        }
        super.visitEnd();
    }

    class JobMethodAdapter extends MethodVisitor {

        public JobMethodAdapter(MethodVisitor mv) {
            super(ASM5, mv);
        }

        @Override
        public void visitInsn(int opcode) {
            if (opcode >= IRETURN && opcode <= RETURN) {
                super.visitVarInsn(ALOAD, 0);
                super.visitFieldInsn(PUTSTATIC, owner, JOB_FIELD_NAME, JOB_FIELD_DESC);
            }
            super.visitInsn(opcode);
        }

        @Override
        public void visitMaxs(int maxStack, int maxLocals) {
            super.visitMaxs(maxStack + 1, maxLocals);
        }
    }
}
